import argparse
import copy
import os
import mindspore
import numpy as np
from pathlib import Path
from mind3d.utils.sim_builder import build_dataset
from mind3d.utils.sim_center_utils import (draw_gaussian, gaussian_radius)
from mind3d.utils.sim_box_np_ops import center_to_corner_box2d
from mind3d.utils.sim_geometry import points_in_convex_polygon_jit
from mind3d.models.build_sim_model import build_model, get_config
from nuscenes.nuscenes import NuScenes

from mindspore import load_checkpoint, ops, Tensor
from mindspore import load_param_into_net
from mindspore import dataset as de
from mindspore import context


def test_simtrack(args):
    
    device_id = int(os.getenv('DEVICE_ID', '1'))
    device_num = int(os.getenv('RANK_SIZE', '1'))
    context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target, device_id=device_id)
    cfg_path = Path(args.config)
    cfg = get_config(cfg_path)

    global voxel_size, downsample, voxel_range, num_classes, size_h, size_w
    voxel_size = np.array(cfg['_voxel_size'])[:2]
    downsample = cfg['assigner']['out_size_factor']
    voxel_range = np.array(cfg['_pc_range'])
    num_classes = sum([t['num_class'] for t in cfg['tasks']])
    size_w, size_h = ((voxel_range[3:5] - voxel_range[:2]) / voxel_size / downsample).astype(np.int32)

    dataset = build_dataset(cfg['data']['val'])
    ms_model = build_model(model_cfg=cfg['model'])

    ds = de.GeneratorDataset(dataset, column_names=cfg['eval_column_names'], shuffle=False)

    ckpt = args.checkpoint
    print(ckpt)
    ms_checkpoint = load_checkpoint(ckpt)
    ms_checkpoint.items()
    load_param_into_net(ms_model, ms_checkpoint)

    expand_dims = ops.ExpandDims()

    for _, data_batch in zip(range(1), enumerate(ds.create_dict_iterator())):
        points = ops.zeros((data_batch['points'].shape[0], 6), mindspore.float32)
        points[:, 1:6] = data_batch['points']
        data_batch['points'] = points
        coor = ops.zeros((data_batch['coordinates'].shape[0], 4), mindspore.float32)
        coor[:, 1:4] = data_batch['coordinates']
        data_batch['coordinates'] = coor
        data_batch['shape'] = data_batch['shape'].view(1, 3).astype('int32')  
        data_batch['ref_from_car'] = expand_dims(data_batch['ref_from_car'], 0)
        data_batch['car_from_global'] = expand_dims(data_batch['car_from_global'], 0)
        data_batch['num_voxels'] = data_batch['num_voxels'].astype('int32')  
        preds=ms_model(data_batch)
    assert isinstance(preds, list), "data type error"
    print("################################")
    print("simtrack test passed!")
    print("################################")



if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="train config file path",
                        default='/mind3d/configs/simtrack/simtrack.yaml')
    parser.add_argument('--device_target', default='GPU')
    parser.add_argument('--device_id',default=1,type=int, help='device id')
    parser.add_argument("--work_dir", help="the dir to save logs and models",
                        default='/mind3d/word_dirs/val/new2_copy/18')
    parser.add_argument(
        "--checkpoint", help="the dir to checkpoint which the model read from",
        default='/mind3d/word_dirs/train/new2_copy/epoch_18.ckpt')
    parser.add_argument("--eval_det", default=False) #不评估检测

    args = parser.parse_args()
    tracking(args)